/*
 * Copyright 2010 Google Inc.
 * 
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not
 * use this file except in compliance with the License. You may obtain a copy of
 * the License at
 * 
 * http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 * License for the specific language governing permissions and limitations under
 * the License.
 */
package com.google.websocket;

import static com.google.websocket.Util.readByte;
import static com.google.websocket.Util.readFully;
import static com.google.websocket.Util.readLength;
import static com.google.websocket.Util.writeLength;

import org.eclipse.jetty.util.Utf8StringBuilder;

import java.io.BufferedInputStream;
import java.io.BufferedOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.nio.charset.Charset;

/**
 * A class representing a WebSocket connection, implemented in a simple and
 * straight-forward manner.
 */
public class SimpleWebSocket {

  // Integer constants are used rather than enums so clients of this library
  // can make use of new (or private-use) frame types before they are defined
  // here.
  public static final int FRAME_CONTINUATION = 0;

  public static final int FRAME_CLOSE = 1;

  public static final int FRAME_PING = 2;

  public static final int FRAME_PONG = 3;

  public static final int FRAME_TEXT = 4;

  public static final int FRAME_BINARY = 5;

  // flags and fields of the first header byte
  private static final int MORE = 0x80;
  private static final int OPCODE_MASK = 0x0f;

  // flags and fields of the second header byte
  private static final int HAS_METADATA = 0x80;

  private static final Charset UTF8 = Charset.forName("UTF-8");

  private final BufferedInputStream bufIn;
  private final BufferedOutputStream bufOut;

  /**
   * Create a WebSocket connection on an already-established connection.  The
   * caller is responsible for completing the connection establishment phase and
   * any extension negotiation before creating this object.
   * 
   * @param in
   * @param out
   */
  public SimpleWebSocket(InputStream in, OutputStream out) {
    bufIn = in instanceof BufferedInputStream ? (BufferedInputStream) in
        : new BufferedInputStream(in);
    bufOut = out instanceof BufferedOutputStream ? (BufferedOutputStream) out
        : new BufferedOutputStream(out);
  }

  /**
   * Close the websocket.
   * 
   * @throws IOException
   */
  public void close() throws IOException {
    sendFrame(false, FRAME_CLOSE, new byte[0]);
    bufOut.close();
  }

  /**
   * Read a text message, processing incoming control frames while doing so.
   * 
   * @return a String containing the text message
   * @throws IOException
   */
  public String readTextMessage() throws IOException {
    Utf8StringBuilder buf = new Utf8StringBuilder();
    int expectedOpcode = FRAME_TEXT;
    boolean moreFrames;
    do {
      int hdrByte = readByte(bufIn);
      moreFrames = (hdrByte & MORE) != 0;
      int opcode = hdrByte & OPCODE_MASK;
      hdrByte = readByte(bufIn);
      assert (hdrByte & HAS_METADATA) == 0;
      long longLength = readLength(hdrByte, bufIn);
      if (longLength > Integer.MAX_VALUE) {
        throw new IOException("Frame size too large");
      }
      int length = (int) longLength;
      byte[] payload = new byte[length];
      readFully(payload, bufIn);
      switch (opcode) {
        case FRAME_CLOSE:
          sendFrame(false, FRAME_CLOSE, new byte[0]);
          bufIn.close();
          throw new IOException("EOF while reading a text message");
        case FRAME_PING:
          sendFrame(false, FRAME_PONG, payload);
          break;
        default:
          if (opcode != expectedOpcode) {
            throw new IOException("Expected opcode " + expectedOpcode + ", found "
                + opcode);
          }
          expectedOpcode = FRAME_CONTINUATION;
          buf.append(payload, 0, length);
          break;
      }
    } while (moreFrames);
    return buf.toString();
  }

  /**
   * Send a text message.
   * 
   * @param message
   * @throws IOException
   */
  public void sendTextMessage(String message)
      throws IOException {
    byte[] bytes = message.getBytes(UTF8);
    sendFrame(false, FRAME_TEXT, bytes);
  }

  /**
   * Send a frame.
   * 
   * @param more true if more frames of this message follow
   * @param opcode frame opcode to use
   * @param payload payload bytes
   * @throws IOException
   */
  // @VisibleForTesting
  void sendFrame(boolean more, int opcode, byte[] payload)
      throws IOException {
    bufOut.write((more ? MORE : 0) | opcode);
    writeLength(bufOut, payload.length, false);
    bufOut.write(payload);
    bufOut.flush();
  }
}
